# from algorithms.generate.stablediffusion.stablediffusion_inferencer import StableDiffusionInferencer
from algorithms import StableDiffusionInferencer

if __name__ == "__main__":
    # 初始化推理器，指定模型名称（如 'CompVis/stable-diffusion-v1-4'）
    # generator = StableDiffusionInferencer("CompVis/stable-diffusion-v1-4", device="cuda")  # 或 "cpu"
    generator = StableDiffusionInferencer("/home/song/data/StableDiffusion/runwayml/stable-diffusion-v1-5", device="cuda")  # 或 "cpu"

    # 输入文本提示
    prompt = "a photograph of an astronaut riding a horse on Mars"

    # 执行文本到图像的生成
    result_images = generator.inference_txt2img(prompt)["images"]

    # 保存生成的图像
    for idx, img in enumerate(result_images):
        generator.save_image(img, f"generated_image_{idx}.png")

    # 显示第一张生成的图像
    generator.show_image(result_images[0])